{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "29ebaee3",
   "metadata": {},
   "source": [
    "# OutlineCompilerFunctionsWithExistingGlobalSymbols 和 MarkCompilerFunctionsAsExtern 外部代码生成\n",
    "\n",
    "用于外部代码生成的两个重要 pass 的测试：\n",
    "1. `OutlineCompilerFunctionsWithExistingGlobalSymbols`：将带有特定编译器标记的函数提取到全局命名空间\n",
    "2. `MarkCompilerFunctionsAsExtern`：将特定编译器的函数标记为外部函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a5fb9638",
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "386a07ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "import tvm.testing\n",
    "import numpy as np\n",
    "\n",
    "def make_const(dtype, shape):\n",
    "    \"\"\"创建指定数据类型和形状的 TVM Relay 常量\"\"\"\n",
    "    return tvm.relay.const(np.random.rand(*shape).astype(dtype))\n",
    "\n",
    "def make_consts(dtype, shapes):\n",
    "    \"\"\"创建多个指定数据类型和形状的 TVM Relay 常量\"\"\"\n",
    "    return [make_const(dtype, shape) for shape in shapes]\n",
    "\n",
    "# 常量表，存储测试中使用的常量\n",
    "metatable = {\n",
    "    \"relay.Constant\": make_consts(\n",
    "        \"float16\",\n",
    "        [\n",
    "            (2304, 768),  # 0: 矩阵乘法的权重\n",
    "            (2304,),      # 1: 偏置向量\n",
    "            (600, 32, 64),  # 2: 批处理矩阵乘法的输入\n",
    "        ],\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2bd5d936",
   "metadata": {},
   "outputs": [],
   "source": [
    "def original_mod():\n",
    "    \"\"\"生成原始的 Relay 模块，包含带有 Cutlass 和 cuBLAS 编译器标记的函数\"\"\"\n",
    "    return tvm.relay.parse(\n",
    "        \"\"\"\n",
    "        #[version = \"0.0.5\"]\n",
    "        def @main(%x0 : Tensor[(1600, 768), float16], %x3 : Tensor[(600, 32, 64), float16]) -> (Tensor[(1600, 2304), float16], Tensor[(600, 32, 32), float16]) {\n",
    "          %0 = fn(%y_0_i0: Tensor[(1600, 768), float16], %y_0_i1: Tensor[(2304, 768), float16], %y_0_i2: Tensor[(2304), float16],\n",
    "                  Inline=1, Compiler=\"cutlass\", global_symbol=\"tvmgen_default_cutlass_main_0\", Primitive=1) -> Tensor[(1600, 2304), float16] {\n",
    "            %4 = fn (%FunctionVar_0_0: Tensor[(1600, 768), float16], %FunctionVar_0_1: Tensor[(2304, 768), float16], %FunctionVar_0_2: Tensor[(2304), float16],\n",
    "                     PartitionedFromPattern=\"nn.dense_add_\", Composite=\"cutlass.dense_bias\") -> Tensor[(1600, 2304), float16] {\n",
    "              %5 = nn.dense(%FunctionVar_0_0, %FunctionVar_0_1, units=2304);\n",
    "              add(%5, %FunctionVar_0_2)\n",
    "            };\n",
    "            %4(%y_0_i0, %y_0_i1, %y_0_i2)\n",
    "          };\n",
    "          %1 = %0(%x0, meta[relay.Constant][0], meta[relay.Constant][1]);\n",
    "          %2 = fn(%y_3_i0: Tensor[(600, 32, 64), float16], %y_3_i1: Tensor[(600, 32, 64), float16],\n",
    "                  Inline=1, Compiler=\"cublas\", global_symbol=\"tvmgen_default_cublas_main_3\", Primitive=1) -> Tensor[(600, 32, 32), float16] {\n",
    "            %6 = fn (%FunctionVar_0_01: Tensor[(600, 32, 64), float16], %FunctionVar_0_11: Tensor[(600, 32, 64), float16],\n",
    "                     PartitionedFromPattern=\"nn.batch_matmul_\", Composite=\"cublas.batch_matmul\") -> Tensor[(600, 32, 32), float16] {\n",
    "              nn.batch_matmul(%FunctionVar_0_01, %FunctionVar_0_11, out_dtype=\"float16\", transpose_b=True)\n",
    "            };\n",
    "            %6(%y_3_i0, %y_3_i1)\n",
    "          };\n",
    "          %3 = %2(%x3, meta[relay.Constant][2]);\n",
    "          (%1, %3)\n",
    "        }\n",
    "        \"\"\",\n",
    "        \"from_string\",\n",
    "        None,\n",
    "        metatable,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "559d83b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>x3: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">34</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cutlass&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cutlass_main_0&quot;</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">47</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">23</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.dense_add_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cutlass.dense_bias&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "      <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>dense(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1, units<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">2304</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "      add(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "    } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">6</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cublas_main_3&quot;</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">31</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">50</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.batch_matmul_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas.batch_matmul&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "      nn<span style=\"color: #A2F; font-weight: bold\">.</span>batch_matmul(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>, transpose_b<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "    } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">16</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">1</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">54</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">23</span>:<span style=\"color: #008000\">12</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x3, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">2</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">23</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ori_mod = original_mod()\n",
    "ori_mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07d0842c",
   "metadata": {},
   "source": [
    "正确地从主函数中提取 Cutlass 编译器标记的函数到全局命名空间："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2054299a",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual_outlined_mod = tvm.relay.transform.OutlineCompilerFunctionsWithExistingGlobalSymbols(\n",
    "    \"cutlass\"\n",
    ")(original_mod())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cd615fa3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>x3: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cublas_main_3&quot;</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">31</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">50</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.batch_matmul_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas.batch_matmul&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "      nn<span style=\"color: #A2F; font-weight: bold\">.</span>batch_matmul(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>, transpose_b<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "    } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">16</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@tvmgen_default_cutlass_main_0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">1</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">54</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">23</span>:<span style=\"color: #008000\">12</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x3, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">2</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">23</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@tvmgen_default_cutlass_main_0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">34</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cutlass&quot;</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cutlass_main_0&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">47</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">23</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.dense_add_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cutlass.dense_bias&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>dense(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1, units<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">2304</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    add(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">6</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "actual_outlined_mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e6ff6d8",
   "metadata": {},
   "source": [
    "验证 let 绑定的 Cutlass 编译器标记函数是否正确地从主函数中提取到全局命名空间："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "21d7d5c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def original_mod_let_bound():\n",
    "    \"\"\"生成 let-bound 形式的原始 Relay 模块，用于测试 let 绑定函数\"\"\"\n",
    "    return tvm.relay.parse(\n",
    "        \"\"\"\n",
    "        #[version = \"0.0.5\"]\n",
    "        def @main(%x0 : Tensor[(1600, 768), float16], %x3 : Tensor[(600, 32, 64), float16]) -> (Tensor[(1600, 2304), float16], Tensor[(600, 32, 32), float16]) {\n",
    "          let %f = fn(%y_0_i0: Tensor[(1600, 768), float16], %y_0_i1: Tensor[(2304, 768), float16], %y_0_i2: Tensor[(2304), float16],\n",
    "                      Inline=1, Compiler=\"cutlass\", global_symbol=\"tvmgen_default_cutlass_main_0\", Primitive=1) -> Tensor[(1600, 2304), float16] {\n",
    "            %4 = fn (%FunctionVar_0_0: Tensor[(1600, 768), float16], %FunctionVar_0_1: Tensor[(2304, 768), float16], %FunctionVar_0_2: Tensor[(2304), float16],\n",
    "                     PartitionedFromPattern=\"nn.dense_add_\", Composite=\"cutlass.dense_bias\") -> Tensor[(1600, 2304), float16] {\n",
    "              %5 = nn.dense(%FunctionVar_0_0, %FunctionVar_0_1, units=2304);\n",
    "              add(%5, %FunctionVar_0_2)\n",
    "            };\n",
    "            %4(%y_0_i0, %y_0_i1, %y_0_i2)\n",
    "          };\n",
    "          %1 = %f(%x0, meta[relay.Constant][0], meta[relay.Constant][1]);\n",
    "          %2 = fn(%y_3_i0: Tensor[(600, 32, 64), float16], %y_3_i1: Tensor[(600, 32, 64), float16],\n",
    "                  Inline=1, Compiler=\"cublas\", global_symbol=\"tvmgen_default_cublas_main_3\", Primitive=1) -> Tensor[(600, 32, 32), float16] {\n",
    "            %6 = fn (%FunctionVar_0_01: Tensor[(600, 32, 64), float16], %FunctionVar_0_11: Tensor[(600, 32, 64), float16],\n",
    "                     PartitionedFromPattern=\"nn.batch_matmul_\", Composite=\"cublas.batch_matmul\") -> Tensor[(600, 32, 32), float16] {\n",
    "              nn.batch_matmul(%FunctionVar_0_01, %FunctionVar_0_11, out_dtype=\"float16\", transpose_b=True)\n",
    "            };\n",
    "            %6(%y_3_i0, %y_3_i1)\n",
    "          };\n",
    "          %3 = %2(%x3, meta[relay.Constant][2]);\n",
    "          (%1, %3)\n",
    "        }\n",
    "        \"\"\",\n",
    "        \"from_string\",\n",
    "        None,\n",
    "        metatable,\n",
    "    )\n",
    "\n",
    "\n",
    "actual_outlined_mod = tvm.relay.transform.OutlineCompilerFunctionsWithExistingGlobalSymbols(\n",
    "    \"cutlass\"\n",
    ")(original_mod_let_bound())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9876ac37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>x3: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cublas_main_3&quot;</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">31</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">50</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.batch_matmul_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas.batch_matmul&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "      nn<span style=\"color: #A2F; font-weight: bold\">.</span>batch_matmul(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>, transpose_b<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">18</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "    } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">20</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">16</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@tvmgen_default_cutlass_main_0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">1</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">54</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">23</span>:<span style=\"color: #008000\">12</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x3, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">2</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">23</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">23</span>:<span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@tvmgen_default_cutlass_main_0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">34</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cutlass&quot;</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cutlass_main_0&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">47</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">23</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.dense_add_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cutlass.dense_bias&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>dense(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1, units<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">2304</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    add(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">8</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">6</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "actual_outlined_mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bea13682",
   "metadata": {},
   "source": [
    "验证外部的 Cutlass 编译器标记函数是否正确地标记为外部函数（Extern=1）："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8567fb8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expected_extern_mod():\n",
    "    \"\"\"将带有 Cutlass 编译器标记的全局函数标记为外部函数（Extern=1）\"\"\"\n",
    "    return tvm.relay.parse(\n",
    "        \"\"\"\n",
    "        #[version = \"0.0.5\"]\n",
    "        def @main(%x0 : Tensor[(1600, 768), float16], %x3 : Tensor[(600, 32, 64), float16]) -> (Tensor[(1600, 2304), float16], Tensor[(600, 32, 32), float16]) {\n",
    "          %1 = @tvmgen_default_cutlass_main_0(%x0, meta[relay.Constant][0], meta[relay.Constant][1]);\n",
    "          %2 = fn(%y_3_i0: Tensor[(600, 32, 64), float16], %y_3_i1: Tensor[(600, 32, 64), float16],\n",
    "                  Inline=1, Compiler=\"cublas\", global_symbol=\"tvmgen_default_cublas_main_3\", Primitive=1) -> Tensor[(600, 32, 32), float16] {\n",
    "            %6 = fn (%FunctionVar_0_01: Tensor[(600, 32, 64), float16], %FunctionVar_0_11: Tensor[(600, 32, 64), float16],\n",
    "                     PartitionedFromPattern=\"nn.batch_matmul_\", Composite=\"cublas.batch_matmul\") -> Tensor[(600, 32, 32), float16] {\n",
    "              nn.batch_matmul(%FunctionVar_0_01, %FunctionVar_0_11, out_dtype=\"float16\", transpose_b=True)\n",
    "            };\n",
    "            %6(%y_3_i0, %y_3_i1)\n",
    "          };\n",
    "          %3 = %2(%x3, meta[relay.Constant][2]);\n",
    "          (%1, %3)\n",
    "        }\n",
    "\n",
    "        def @tvmgen_default_cutlass_main_0(%y_0_i0: Tensor[(1600, 768), float16], %y_0_i1: Tensor[(2304, 768), float16], %y_0_i2: Tensor[(2304), float16],\n",
    "                  Extern=1) -> Tensor[(1600, 2304), float16] {\n",
    "          %4 = fn (%FunctionVar_0_0: Tensor[(1600, 768), float16], %FunctionVar_0_1: Tensor[(2304, 768), float16], %FunctionVar_0_2: Tensor[(2304), float16],\n",
    "                   PartitionedFromPattern=\"nn.dense_add_\", Composite=\"cutlass.dense_bias\") -> Tensor[(1600, 2304), float16] {\n",
    "            %5 = nn.dense(%FunctionVar_0_0, %FunctionVar_0_1, units=2304);\n",
    "            add(%5, %FunctionVar_0_2)\n",
    "          };\n",
    "          %4(%y_0_i0, %y_0_i1, %y_0_i2)\n",
    "        }\n",
    "        \"\"\",\n",
    "        \"from_string\",\n",
    "        None,\n",
    "        metatable,\n",
    "    )\n",
    "\n",
    "actual_extern_mod = tvm.relay.transform.MarkCompilerFunctionsAsExtern(\"cutlass\")(\n",
    "    expected_extern_mod()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "79981adb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">47</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>x3: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cublas_main_3&quot;</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">31</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">50</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.batch_matmul_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas.batch_matmul&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "      nn<span style=\"color: #A2F; font-weight: bold\">.</span>batch_matmul(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>, transpose_b<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "    } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">7</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@tvmgen_default_cutlass_main_0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">57</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">1</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">82</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">14</span>:<span style=\"color: #008000\">12</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x3, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">2</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">14</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@tvmgen_default_cutlass_main_0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">24</span>:<span style=\"color: #008000\">14</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">24</span>:<span style=\"color: #008000\">23</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">24</span>:<span style=\"color: #008000\">32</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Extern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">21</span>:<span style=\"color: #008000\">27</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">21</span>:<span style=\"color: #008000\">45</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">2304</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">21</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.dense_add_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cutlass.dense_bias&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>dense(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1, units<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">2304</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">17</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    add(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">21</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16], Tensor[(<span style=\"color: #008000\">2304</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">24</span>:<span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i1, <span style=\"color: #A2F; font-weight: bold\">%</span>y_0_i2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">19</span>:<span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "actual_extern_mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c5ca5c9",
   "metadata": {},
   "source": [
    "验证指定的全局函数是否正确地内联回主函数体中："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c009a1f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expected_outlined_mod():\n",
    "    \"\"\"将带有 Cutlass 编译器标记的函数从主函数体中提取到全局命名空间\"\"\"\n",
    "    return tvm.relay.parse(\n",
    "        \"\"\"\n",
    "        #[version = \"0.0.5\"]\n",
    "        def @main(%x0 : Tensor[(1600, 768), float16], %x3 : Tensor[(600, 32, 64), float16]) -> (Tensor[(1600, 2304), float16], Tensor[(600, 32, 32), float16]) {\n",
    "          %1 = @tvmgen_default_cutlass_main_0(%x0, meta[relay.Constant][0], meta[relay.Constant][1]);\n",
    "          %2 = fn(%y_3_i0: Tensor[(600, 32, 64), float16], %y_3_i1: Tensor[(600, 32, 64), float16],\n",
    "                  Inline=1, Compiler=\"cublas\", global_symbol=\"tvmgen_default_cublas_main_3\", Primitive=1) -> Tensor[(600, 32, 32), float16] {\n",
    "            %6 = fn (%FunctionVar_0_01: Tensor[(600, 32, 64), float16], %FunctionVar_0_11: Tensor[(600, 32, 64), float16],\n",
    "                     PartitionedFromPattern=\"nn.batch_matmul_\", Composite=\"cublas.batch_matmul\") -> Tensor[(600, 32, 32), float16] {\n",
    "              nn.batch_matmul(%FunctionVar_0_01, %FunctionVar_0_11, out_dtype=\"float16\", transpose_b=True)\n",
    "            };\n",
    "            %6(%y_3_i0, %y_3_i1)\n",
    "          };\n",
    "          %3 = %2(%x3, meta[relay.Constant][2]);\n",
    "          (%1, %3)\n",
    "        }\n",
    "\n",
    "        def @tvmgen_default_cutlass_main_0(%y_0_i0: Tensor[(1600, 768), float16], %y_0_i1: Tensor[(2304, 768), float16], %y_0_i2: Tensor[(2304), float16],\n",
    "                  Inline=1, Compiler=\"cutlass\", global_symbol=\"tvmgen_default_cutlass_main_0\", Primitive=1) -> Tensor[(1600, 2304), float16] {\n",
    "          %4 = fn (%FunctionVar_0_0: Tensor[(1600, 768), float16], %FunctionVar_0_1: Tensor[(2304, 768), float16], %FunctionVar_0_2: Tensor[(2304), float16],\n",
    "                   PartitionedFromPattern=\"nn.dense_add_\", Composite=\"cutlass.dense_bias\") -> Tensor[(1600, 2304), float16] {\n",
    "            %5 = nn.dense(%FunctionVar_0_0, %FunctionVar_0_1, units=2304);\n",
    "            add(%5, %FunctionVar_0_2)\n",
    "          };\n",
    "          %4(%y_0_i0, %y_0_i1, %y_0_i2)\n",
    "        }\n",
    "        \"\"\",\n",
    "        \"from_string\",\n",
    "        None,\n",
    "        metatable,\n",
    "    )\n",
    "mod = expected_outlined_mod()\n",
    "gv = mod.get_global_var(\"tvmgen_default_cutlass_main_0\")\n",
    "actual_inlined_mod = tvm.relay.transform.InlineCompilerFunctionsBoundTo([gv])(mod)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fdbd6a1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">47</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>x3: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">19</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>dense(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>, <span style=\"color: #008000\">768</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">57</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, units<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">2304</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">22</span>:<span style=\"color: #008000\">17</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">25</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_cublas_main_3&quot;</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">31</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11: Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">50</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.batch_matmul_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;cublas.batch_matmul&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] {\n",
       "      nn<span style=\"color: #A2F; font-weight: bold\">.</span>batch_matmul(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_11, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>, transpose_b<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">9</span>:<span style=\"color: #008000\">15</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "    } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">11</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i0, <span style=\"color: #A2F; font-weight: bold\">%</span>y_3_i1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">7</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> add(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">1</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">82</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">21</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x3, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">2</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">64</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">13</span>:<span style=\"color: #008000\">29</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">14</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1600</span>, <span style=\"color: #008000\">2304</span>), float16], Tensor[(<span style=\"color: #008000\">600</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float16]) span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "actual_inlined_mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd05738d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
